import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/index.js";
import { MAX_BATCH_RETRIES, INITIAL_RETRY_DELAY_MS } from "../constants/index.js";
import { formatEmbeddingError, withValidationErrorHandling } from "../utils/validation-helpers.js";
import { InstructionManager, DocumentType } from "../utils/instruction-manager.js";

export class OllamaEmbedder implements IEmbedder {
  private readonly baseUrl: string;
  private readonly defaultModelId: string;

  constructor(options: { baseUrl: string; modelId?: string }) {
    this.baseUrl = options.baseUrl.replace(/\/$/, '');
    this.defaultModelId = options.modelId || "nomic-embed-text";
  }

  async createEmbeddings(texts: string[], docType: DocumentType, task: 'retrieval_query' | 'retrieval_document'): Promise<EmbeddingResponse> {
    const embeddings = await Promise.all(
      texts.map(text => {
        const instructedText = InstructionManager.getInstruction(task, docType, text);
        return this._embedWithRetries(instructedText, this.defaultModelId);
      })
    );
    return { embeddings };
  }

  private async _embedWithRetries(text: string, model: string): Promise<number[]> {
    for (let attempts = 0; attempts < MAX_BATCH_RETRIES; attempts++) {
      try {
        const response = await fetch(`${this.baseUrl}/api/embeddings`, {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({ model, prompt: text }),
        });

        if (!response.ok) throw new Error(`Ollama API error: ${response.status} ${response.statusText}`);
        const data = await response.json();
        if (!data.embedding || !Array.isArray(data.embedding)) throw new Error("Invalid response format from Ollama API");
        return data.embedding;
      } catch (error: any) {
        if (attempts < MAX_BATCH_RETRIES - 1) {
          const delayMs = INITIAL_RETRY_DELAY_MS * Math.pow(2, attempts);
          await new Promise((resolve) => setTimeout(resolve, delayMs));
        } else {
          throw formatEmbeddingError(error, MAX_BATCH_RETRIES);
        }
      }
    }
    throw new Error(`Failed to create embedding after ${MAX_BATCH_RETRIES} attempts`);
  }

  async validateConfiguration(): Promise<{ valid: boolean; error?: string }> {
    return withValidationErrorHandling(async () => {
      const response = await fetch(`${this.baseUrl}/api/tags`);
      if (!response.ok) throw new Error(`Cannot connect to Ollama server: ${response.status} ${response.statusText}`);
      const data = await response.json();
      const modelExists = (data.models || []).some((m: any) => m.name.startsWith(this.defaultModelId));
      if (!modelExists) return { valid: false, error: `Model '${this.defaultModelId}' not found.` };
      const testEmbedding = await this._embedWithRetries("test", this.defaultModelId);
      if (!testEmbedding || testEmbedding.length === 0) return { valid: false, error: "Embedding test returned empty result" };
      return { valid: true };
    }, "ollama");
  }

  get embedderInfo(): EmbedderInfo {
    return { name: "ollama" };
  }
}
